#include "kernel_operator.h"
using namespace AscendC;

extern "C" __global__ __aicore__ void add_kernel(__gm__ float* x, __gm__ float* y, __gm__ float* z, const uint32_t sz)
{
  const int buf_sz = 2, block_sz = 1024, tile_sz = 128;
  TPipe pipe;
  TQue<QuePosition::VECIN, buf_sz> inQueueX, inQueueY;
  TQue<QuePosition::VECOUT, buf_sz> outQueueZ;
  GlobalTensor<float> xGm;
  GlobalTensor<float> yGm;
  GlobalTensor<float> zGm;
  xGm.SetGlobalBuffer(x + block_sz * GetBlockIdx(), block_sz);
  yGm.SetGlobalBuffer(y + block_sz * GetBlockIdx(), block_sz);
  zGm.SetGlobalBuffer(z + block_sz * GetBlockIdx(), block_sz);
  // printf("%f\n", xGm.GetValue(0));
  pipe.InitBuffer(inQueueX, buf_sz, tile_sz * sizeof(float));
  pipe.InitBuffer(inQueueY, buf_sz, tile_sz * sizeof(float));
  pipe.InitBuffer(outQueueZ, buf_sz, tile_sz * sizeof(float));
  int32_t block_tsz =  block_sz/tile_sz;
  for (int32_t i = 0; i < block_tsz; i++) {
    LocalTensor<float> xLocal = inQueueX.AllocTensor<float>(), yLocal = inQueueY.AllocTensor<float>();
    DataCopy(xLocal, xGm[i*tile_sz], tile_sz);
    DataCopy(yLocal, yGm[i*tile_sz], tile_sz);
    inQueueX.EnQue(xLocal);
    inQueueY.EnQue(yLocal);
    xLocal = inQueueX.DeQue<float>();
    yLocal = inQueueY.DeQue<float>();
    LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();
    AscendC::Add(zLocal, xLocal, yLocal, tile_sz);
    outQueueZ.EnQue<float>(zLocal);
    inQueueX.FreeTensor(xLocal);
    inQueueY.FreeTensor(yLocal);
    zLocal = outQueueZ.DeQue<float>();
    DataCopy(zGm[i * tile_sz], zLocal, tile_sz);
    outQueueZ.FreeTensor(zLocal);
  }
  // printf("z %f\n", z[0]);
}

void AddCus(uint32_t blockdim, void *stream, float * x, float *y, float *z, uint32_t sz)
{
  add_kernel<<<blockdim, nullptr, stream>>>(x, y, z, sz);
}
